import numpy as np
import cv2
import json
import triton_python_backend_utils as pb_utils


class TritonPythonModel:

    def initialize(self, args):
        self.logger = pb_utils.Logger

        # parse model configs
        self.model_config = json.loads(args['model_config'])

        output_config = pb_utils.get_output_config_by_name(self.model_config, "output_tensors")
        self.output_dims = output_config['dims']
        self.output_h = self.output_dims[1]
        self.output_w = self.output_dims[2]
        # self.logger.log_info('output_dims: {}'.format(self.output_dims))
        # self.logger.log_info('output_h: {}'.format(self.output_h))
        # self.logger.log_info('output_w: {}'.format(self.output_w))

        # parse parameters
        self.parameters = self.model_config['parameters']
        self.model_name = self.model_config['name']

        self.torch_device   = self.parameters['torch_device']['string_value']
        self.logger.log_info('torch_device: {}'.format(self.torch_device)) #若正常运行，前处理没有torch操作，因此有无该参数无所谓了

        self.logger.log_info('{} started'.format(self.model_name))

    
    def execute(self, requests):
        print('running {}'.format(self.model_name))
        # print(type(requests)) #List
        responses = []
        for request in requests:
            # print(type(request))
            images = pb_utils.get_input_tensor_by_name(request, 'images').as_numpy()#nhwc
            # print(type(images))
            # print(images.shape)

            output_tensors = []
            org_img_hw = []
            for i in range(images.shape[0]):
                image = images[i, :, :, :]#hwc
                org_img_hw.append(np.array(image.shape[:2], dtype=np.int32))
                output_tensor = self.prepare_input(image)#1chw
                output_tensors.append(output_tensor)

            output_tensors = np.vstack(output_tensors)
            org_img_hw = np.vstack(org_img_hw)

            inference_response = pb_utils.InferenceResponse(
                output_tensors=[
                    pb_utils.Tensor(
                        "output_tensors",
                        output_tensors,
                    ),
                    pb_utils.Tensor(
                        "org_img_hw",
                        org_img_hw,
                    ),
                ]
            )
            responses.append(inference_response)

        return responses


    def finalize(self):
        self.logger.log_info('Finalizing {}'.format(self.model_name))


    def prepare_input(self, image: np.ndarray) -> np.ndarray:
        # https://gitee.com/luo_zhi_cheng/sam2/blob/main/sam2/utils/transforms.py
        # https://gitee.com/luo_zhi_cheng/ONNX-SAM2-Segment-Anything/blob/main/sam2/sam2.py 的 prepare_input
        # image -> hwc from cv.imread
        
        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Resize input image
        input_img = cv2.resize(input_img, (self.output_w, self.output_h))

        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        input_img = (input_img / 255.0 - mean) / std
        input_img = input_img.transpose(2, 0, 1)
        input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)

        return input_tensor
